# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph
from experimental.prophet.ops.util import (to_tuple, to_data_type_proto)
from experimental.prophet.ops.constant import constant

_null_var_name_map = {
    to_data_type_proto(int): 'null_int',
    to_data_type_proto(float): 'null_double',
    to_data_type_proto(bool): 'null_bool'
}


def null(dtype, shape=None):
  assert dtype is not None
  var_name = _null_var_name_map[to_data_type_proto(dtype)]
  try:
    return graph.get_named_variable(var_name)
  except KeyError:
    pass

  var = constant(None, dtype=dtype, node_name=var_name, namespace='')
  return var.named_ref(var_name, namespace='')


def is_null_op(values):
  values = [graph.as_variable(v) for v in to_tuple(values)]
  outs = []
  for v in values:
    node = graph.create_node('IsNull', [v], attrs={'dtype': to_data_type_proto(v.dtype)})
    outs.append(node.outputs[0])
  return outs
